"""
# -*- coding: utf-8 -*-
# @Time    : 2023/5/24 15:10
# @Author  : 王摇摆
# @FileName: Model_Manual.py
# @Software: PyCharm
# @Blog    ：https://blog.csdn.net/weixin_44943389?type=blog
"""
import numpy as np
from sklearn.tree import DecisionTreeRegressor


class gbdtc:
    """
    梯度提升树二分类算法
    """

    def __init__(self, n_estimators=100, learning_rate=0.1):
        # 梯度提升树弱学习器数量
        self.n_estimators = n_estimators
        # 学习速率
        self.learning_rate = learning_rate
        print('人工GBDT二分类树已初始化完毕！')

    def fit(self, X, y):
        """
        梯度提升树二分类算法拟合
        """
        # 标签类
        self.y_classes = np.unique(y)
        # 标签类数量
        self.n_classes = len(self.y_classes)
        # 标签的平均值
        y_avg = np.average(y)
        # 初始化H0
        self.H0 = np.log((1 + y_avg) / (1 - y_avg)) / 2
        # 初始化预测值
        H = np.ones(X.shape[0]) * self.H0
        # 估计器数组
        estimators = []
        # 叶子结点取值数组
        gammas = []
        for k in range(self.n_estimators):
            # 计算 y_hat
            y_hat = 2 * np.multiply(y, 1 / (1 + np.exp(2 * np.multiply(y, H))))
            # 初始化决策回归树估计器
            estimator = DecisionTreeRegressor(max_depth=3, criterion="friedman_mse")
            # 将 y_hat 当作标签值拟合训练集
            estimator.fit(X, y_hat)
            # 计算训练集在当前决策回归树的叶子结点
            leaf_ids = estimator.apply(X)
            # 每个叶子结点下包含的训练数据序号
            node_ids_dict = self.get_leaf_nodes(leaf_ids)
            # 叶子结点取值字典表
            gamma_dict = {}
            # 计算叶子结点取值
            for leaf_id, node_ids in node_ids_dict.items():
                # 当前叶子结点包含的 y_hat
                y_hat_sub = y_hat[node_ids]
                y_hat_sub_abs = np.abs(y_hat_sub)
                # 计算叶子结点取值
                gamma = np.sum(y_hat_sub) / np.sum(np.multiply(y_hat_sub_abs, 2 - y_hat_sub_abs))
                gamma_dict[leaf_id] = gamma
                # 更新预测值
                H[node_ids] += self.learning_rate * gamma
            estimators.append(estimator)
            gammas.append(gamma_dict)
        self.estimators = estimators
        self.gammas = gammas

    def predict(self, X):
        """
        梯度提升树二分类算法预测
        """
        # 初始化预测值
        H = np.ones(X.shape[0]) * self.H0
        # 遍历估计器
        for k in range(self.n_estimators):
            estimator = self.estimators[k]
            # 计算在当前决策回归树的叶子结点
            leaf_ids = estimator.apply(X)
            # 每个叶子结点下包含的数据序号
            node_ids_dict = self.get_leaf_nodes(leaf_ids)
            # 叶子结点取值字典表
            gamma_dict = self.gammas[k]
            # 计算预测值
            for leaf_id, node_ids in node_ids_dict.items():
                gamma = gamma_dict[leaf_id]
                H[node_ids] += self.learning_rate * gamma
        # 计算概率
        probs = np.zeros((X.shape[0], self.n_classes))
        probs[:, 0] = 1 / (1 + np.exp(2 * H))
        probs[:, 1] = 1 / (1 + np.exp(-2 * H))
        return self.y_classes.take(np.argmax(probs, axis=1), axis=0)

    def get_leaf_nodes(self, leaf_ids):
        """
        每个叶子结点下包含的数据序号
        """
        node_ids_dict = {}
        for j in range(len(leaf_ids)):
            leaf_id = leaf_ids[j]
            node_ids = node_ids_dict.setdefault(leaf_id, [])
            node_ids.append(j)
        return node_ids_dict
